#!/usr/bin/env python3
"""
IndexTTS2 环境检查脚本
用于验证部署环境是否正确配置
"""

import os
import sys
import platform
import subprocess
from pathlib import Path
import importlib

def check_python():
    """检查Python版本"""
    print("🐍 Python环境检查:")
    print(f"   版本: {sys.version}")
    print(f"   路径: {sys.executable}")
    
    # 检查Python版本
    version = sys.version_info
    if version.major == 3 and version.minor >= 9:
        print("   ✅ Python版本符合要求 (3.9+)")
        return True
    else:
        print("   ❌ Python版本过低，需要3.9+")
        return False

def check_system():
    """检查系统信息"""
    print("\n💻 系统环境检查:")
    print(f"   操作系统: {platform.system()} {platform.release()}")
    print(f"   架构: {platform.machine()}")
    print(f"   处理器: {platform.processor()}")
    
    if platform.system() == "Windows":
        print("   ✅ Windows系统")
        return True
    else:
        print("   ⚠️  非Windows系统，部分功能可能不可用")
        return False

def check_cuda():
    """检查CUDA环境"""
    print("\n🚀 CUDA环境检查:")
    
    try:
        # 检查nvidia-smi
        result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
        if result.returncode == 0:
            print("   ✅ NVIDIA驱动已安装")
            
            # 尝试获取CUDA版本
            try:
                result = subprocess.run(['nvcc', '--version'], capture_output=True, text=True)
                if result.returncode == 0:
                    print("   ✅ CUDA Toolkit已安装")
                else:
                    print("   ⚠️  CUDA Toolkit未安装或不在PATH中")
            except FileNotFoundError:
                print("   ⚠️  nvcc命令未找到")
            
            return True
        else:
            print("   ❌ NVIDIA驱动未安装")
            return False
    except FileNotFoundError:
        print("   ❌ nvidia-smi命令未找到")
        return False

def check_packages():
    """检查Python包"""
    print("\n📦 Python包检查:")
    
    required_packages = {
        'torch': 'PyTorch',
        'torchaudio': 'TorchAudio', 
        'librosa': 'Librosa',
        'transformers': 'Transformers',
        'gradio': 'Gradio',
        'fastapi': 'FastAPI',
        'uvicorn': 'Uvicorn',
        'pywin32': 'PyWin32',
        'psutil': 'PSUtil',
        'omegaconf': 'OmegaConf',
        'safetensors': 'SafeTensors'
    }
    
    all_packages_ok = True
    
    for package, description in required_packages.items():
        try:
            module = importlib.import_module(package)
            version = getattr(module, '__version__', 'Unknown')
            print(f"   ✅ {description}: {version}")
        except ImportError:
            print(f"   ❌ {description}: 未安装")
            all_packages_ok = False
    
    return all_packages_ok

def check_torch_gpu():
    """检查PyTorch GPU支持"""
    print("\n🔥 PyTorch GPU检查:")
    
    try:
        import torch
        print(f"   PyTorch版本: {torch.__version__}")
        
        if torch.cuda.is_available():
            print("   ✅ CUDA可用")
            print(f"   GPU数量: {torch.cuda.device_count()}")
            for i in range(torch.cuda.device_count()):
                print(f"   GPU {i}: {torch.cuda.get_device_name(i)}")
            return True
        else:
            print("   ❌ CUDA不可用，将使用CPU模式")
            return False
    except ImportError:
        print("   ❌ PyTorch未安装")
        return False

def check_model_files():
    """检查模型文件"""
    print("\n🎵 模型文件检查:")
    
    model_files = [
        'checkpoints/config.yaml',
        'checkpoints/gpt.pth',
        'checkpoints/s2mel.pth', 
        'checkpoints/bpe.model',
        'checkpoints/wav2vec2bert_stats.pt'
    ]
    
    all_files_ok = True
    
    for file_path in model_files:
        if Path(file_path).exists():
            size = Path(file_path).stat().st_size / (1024*1024)  # MB
            print(f"   ✅ {file_path}: {size:.1f}MB")
        else:
            print(f"   ❌ {file_path}: 文件不存在")
            all_files_ok = False
    
    return all_files_ok

def check_directories():
    """检查目录结构"""
    print("\n📁 目录结构检查:")
    
    required_dirs = [
        'checkpoints',
        'indextts', 
        'examples',
        'outputs',
        'uploads'
    ]
    
    all_dirs_ok = True
    
    for dir_path in required_dirs:
        if Path(dir_path).exists():
            print(f"   ✅ {dir_path}/")
        else:
            print(f"   ❌ {dir_path}/ 目录不存在")
            if dir_path in ['outputs', 'uploads']:
                print(f"      正在创建 {dir_path}/ 目录...")
                Path(dir_path).mkdir(exist_ok=True)
                print(f"      ✅ {dir_path}/ 目录创建成功")
            else:
                all_dirs_ok = False
    
    return all_dirs_ok

def check_windows_service():
    """检查Windows服务"""
    print("\n🔧 Windows服务检查:")
    
    if not Path('windows_service.py').exists():
        print("   ❌ windows_service.py 文件不存在")
        return False
    
    try:
        # 检查服务是否已安装
        result = subprocess.run(['sc', 'query', 'IndexTTSService'], 
                              capture_output=True, text=True)
        if result.returncode == 0:
            print("   ✅ IndexTTSService 服务已安装")
            
            # 检查服务状态
            if "RUNNING" in result.stdout:
                print("   ✅ 服务正在运行")
            elif "STOPPED" in result.stdout:
                print("   ⚠️  服务已停止")
            else:
                print("   ⚠️  服务状态未知")
            return True
        else:
            print("   ❌ IndexTTSService 服务未安装")
            return False
    except Exception as e:
        print(f"   ❌ 无法检查服务状态: {e}")
        return False

def main():
    """主检查函数"""
    print("=" * 50)
    print("       IndexTTS2 环境检查报告")
    print("=" * 50)
    
    results = []
    
    # 执行各项检查
    results.append(("Python环境", check_python()))
    results.append(("系统环境", check_system()))
    results.append(("CUDA环境", check_cuda()))
    results.append(("Python包", check_packages()))
    results.append(("PyTorch GPU", check_torch_gpu()))
    results.append(("模型文件", check_model_files()))
    results.append(("目录结构", check_directories()))
    results.append(("Windows服务", check_windows_service()))
    
    # 生成总结报告
    print("\n" + "=" * 50)
    print("           检查结果总结")
    print("=" * 50)
    
    passed = 0
    total = len(results)
    
    for item, result in results:
        status = "✅ 通过" if result else "❌ 失败"
        print(f"{item:12} : {status}")
        if result:
            passed += 1
    
    print("-" * 50)
    print(f"通过率: {passed}/{total} ({passed/total*100:.0f}%)")
    
    if passed == total:
        print("🎉 所有检查通过！环境配置正确。")
        print("💡 您可以启动服务了:")
        print("   net start IndexTTSService")
        print("   或访问: http://127.0.0.1:7863")
    else:
        print("⚠️  存在配置问题，请根据上述提示进行修复。")
        print("💡 常见解决方案:")
        print("   1. 运行 auto_deploy.bat 重新部署")
        print("   2. 检查 DEPLOYMENT_GUIDE.md 详细说明")
        print("   3. 手动安装缺失的依赖包")

if __name__ == "__main__":
    main()
